// arch/riscv/kernel/vm.c

#include <types.h>
#include <defs.h>
#include "string.h"
#include "mm.h"
#include "printk.h"

extern uint64 _stext;
extern uint64 _etext;
extern uint64 _srodata;
extern uint64 _erodata;
extern uint64 _sdata;
extern uint64 _edata;
extern uint64 _sbss;
extern uint64 _ebss;

void create_mapping(uint64 *pgtbl, uint64 va, uint64 pa, uint64 sz, int perm);

struct pgtblEntry_bit{
    unsigned short V: 1;
    unsigned short R: 1;
    unsigned short W: 1;
    unsigned short X: 1;
    unsigned short U: 1;
    unsigned short G: 1;
    unsigned short A: 1;
    unsigned short D: 1;
    unsigned short RSW: 2;
    unsigned long ppn: 26;
    unsigned int reserved: 28;
};

/* early_pgtbl: 用于 setup_vm 进行 1GB 的 映射。 */
uint64 early_pgtbl[512] __attribute__((__aligned__(0x1000)));

struct va_bit
{
    uint64 offset: 30;
    uint64 vpn: 9;
    uint64 high_bit : 25;
};


struct pa_bit
{
    unsigned int offset: 30;
    unsigned long ppn: 26;
    unsigned short high_bit : 8;
};

struct va_sv39
{
    unsigned int offset: 12;
    unsigned short vpn0: 9;
    unsigned short vpn1: 9;
    unsigned short vpn2: 9;
    unsigned int high_bot: 25;
};

struct pa_sv39
{
    unsigned int offset: 12;
    unsigned long ppn: 44;
    unsigned short high_bit: 8;
};

void setup_vm(void) {

    memset(early_pgtbl, 0x0, PGSIZE);

    //1. 由于是进行 1GB 的映射 这里不需要使用多级页表 
    //2. 将 va 的 64bit 作为如下划分： | high bit | 9 bit | 30 bit |
    //high bit 可以忽略
    //中间9 bit 作为 early_pgtbl 的 index
    struct pa_sv39 pa;
    struct va_bit va;
    uint64 i = PHY_START;

    pa = *(struct pa_sv39*)(&i);
    printk("pa.ppn: %x, pa.offset: %x\n", pa.ppn, pa.offset);
    int pa_index = pa.ppn;
    uint64 va_index1 = i;
    uint64 va_index2 = i + PA2VA_OFFSET;
    struct va_bit va1;
    struct va_bit va2;
    va1 = *(struct va_bit*)(&va_index1);
    va2 = *(struct va_bit*)(&va_index2);
    printk("%x %x ", va1.vpn, va2.vpn);

    struct pgtblEntry_bit *entry1 = (struct pgtblEntry_bit*)&early_pgtbl[va1.vpn];
    struct pgtblEntry_bit *entry2 = (struct pgtblEntry_bit*)&early_pgtbl[va2.vpn];
    entry1->ppn = pa_index;
    entry2->ppn = pa_index;
    
    // 00000002 00000180

    //低 30 bit 作为 页内偏移 这里注意到 30 = 9 + 9 + 12， 即我们只使用根页表， 根页表的每个 entry 都对应 1GB 的区域。 
    //3. Page Table Entry 的权限 V | R | W | X 位设置为 1
    
    entry1->V = 1;
    // printk("entry1: %x, entry2: %x\n", entry1, entry2);
    entry1->R = 1;
    entry1->W = 1;
    entry1->X = 1;

    entry2->V = 1;
    entry2->R = 1;
    entry2->W = 1;
    entry2->X = 1;
    printk("entry1: %x, entry2: %x\n", early_pgtbl[va1.vpn], early_pgtbl[va2.vpn]);
}

// arch/riscv/kernel/vm.c 

/* swapper_pg_dir: kernel pagetable 根目录， 在 setup_vm_final 进行映射。 */
uint64  swapper_pg_dir[512] __attribute__((__aligned__(0x1000)));



void setup_vm_final(void) {
    //printk("ENTER setup_vm_final\n");
    memset(swapper_pg_dir, 0x0, PGSIZE);

    // No OpenSBI mapping required
    uint64 stext = (uint64)&_stext;
    uint64 etext = (uint64)&_etext;
    uint64 srodata = (uint64)&_srodata;
    uint64 erodata = (uint64)&_erodata;
    uint64 sdata = (uint64)&_sdata;
    // mapping kernel text X|-|R|V
    create_mapping(swapper_pg_dir, stext-PA2VA_OFFSET, stext-PA2VA_OFFSET,  
                    etext-stext, 0x7);
    create_mapping(swapper_pg_dir, stext, stext-PA2VA_OFFSET,  
                    etext-stext, 0x5);
    //printk("create mapping finished.\n");

    // mapping kernel rodata -|-|R|V
    create_mapping(swapper_pg_dir, srodata, srodata-PA2VA_OFFSET, 
                    erodata-srodata, 0x1);
    //printk("create mapping finished.\n");

    
    // mapping other memory -|W|R|V
    uint64 other_mem_size = PHY_SIZE - (erodata - srodata) - (etext - stext) - OPENSBI_SIZE;
    create_mapping(swapper_pg_dir, sdata, sdata-PA2VA_OFFSET, 
                    other_mem_size, 0x3);
    //printk("create mapping finished.\n");

    
    // set satp with swapper_pg_dir
    uint64 PPN = ((uint64)(swapper_pg_dir) - PA2VA_OFFSET) >> 12;
    uint64 mode = 8;
    uint64 satp_value = (mode<<60)|(PPN&0xfffffffffff);
    csr_write(satp,satp_value);
    //printk("Set satp finished.\n");


    // flush TLB
    asm volatile("sfence.vma zero, zero");
    printk("TLB flushed in setup_vm_final\n");
    return;
}


/* 创建多级页表映射关系 */
void create_mapping(uint64 *pgtbl, uint64 va, uint64 pa, uint64 sz, int perm) {
    /*
    pgtbl 为根页表的基地址
    va, pa 为需要映射的虚拟地址、物理地址
    sz 为映射的大小
    perm 为映射的读写权限

    创建多级页表的时候可以使用 kalloc() 来获取一页作为页表目录
    可以使用 V bit 来判断页表项是否存在

    */
    // printk("M\n");
    uint64 VA = va, PA = pa;
    for(; VA < va+sz; VA+=PGSIZE, PA+=PGSIZE) {
        // printk("\nNew Time:\n");
        uint64* pgtbl_ = pgtbl;
        // printk("address of head pgtbl_: %016lx",pgtbl_);
        //Fetch vpn[2], vpn[1] and vpn[0]
        uint64 vpn[3];
        vpn[2] = (uint64)((VA>>30) & (0x1ff));
        vpn[1] = (uint64)((VA>>21) & (0x1ff));
        vpn[0] = (uint64)((VA>>12) & (0x1ff));
        // printk("Complete setting vpn\n");
        //On the layer2
        if(pgtbl_[vpn[2]] & 0x1){
            // printk("Layer2 is valid\n");
            //layer1 has already been created
            pgtbl_ = (uint64)(((pgtbl_[vpn[2]] >> 10) << 12) + PA2VA_OFFSET);
            // printk("address of new pgtbl_: %016lx\n",pgtbl_);
        }
        else{
            // printk("Layer2 is not valid\n");
            //the layer1 should be created first
            uint64* pgtbl_layer1_address = (uint64*)kalloc();
            //Write that address into layer2
            pgtbl_[vpn[2]] = ((((uint64)(pgtbl_layer1_address)-PA2VA_OFFSET)>>12)<<10) | 0x1;
            //Go from layer2 to layer1
            pgtbl_ = pgtbl_layer1_address;
        }
        // printk("layer2 finished\n");
        //On the layer1
        if(pgtbl_[vpn[1]] & 0x1){
            // printk("layer1 isValid\n");
            //layer0 has already been created
            pgtbl_ = (uint64*)(((pgtbl_[vpn[1]] >> 10) << 12) + PA2VA_OFFSET);
        }
        else{
            // printk("layer1 notValid\n");
            //the layer0 should be created first
            uint64* pgtbl_layer0_address = (uint64*)kalloc();
            //Write that address into layer1
            pgtbl_[vpn[1]] = ((((uint64)(pgtbl_layer0_address)-PA2VA_OFFSET)>>12)<<10) | 0x1;
            //MOve from layer1 to layer0
            pgtbl_ = pgtbl_layer0_address;
        }
        // printk("layer1 finished\n");
        //On the layer0
        //The address should be modified to pa
        //First get the attribute PPN and OFFSET from pa
        uint64 PPN = (PA>>12) & 0xfffffffffff;
        uint64 offset = PA & 0xfff;
        //Second write PPN into the layer0
        
        pgtbl_[vpn[0]] = ((uint64)PPN<<10) | (perm<<1) | 0x1;
        // printk("pgtbl_[vpn[0]]: 0x%016lx\n",pgtbl_[vpn[0]]);
        // printk("layer0 finished\n");
    }
}
